package com.kh.rationaldrugreview.service;

import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;

/**
 * Enhanced Drug Manual Parser using Spring AI
 * Provides structured prompting and response handling
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class DrugManualAIParser {
    
    private final ChatClient chatClient;
    
    private static final String SYSTEM_PROMPT = """
        You are a pharmaceutical expert AI that analyzes drug manuals and extracts structured information.
        Follow the exact format specified for entity extraction and relationship building.
        Always use the delimiter format specified and ensure all entities have complete attributes.
        """;
    
    private static final String EXTRACTION_TEMPLATE = """
        Parse the following drug manual and extract entities and relationships.
        
        Entity Format: ("entity"<|><NAME><|><TYPE><|>"attributes": <Attribute1>=<Value1>;<Attribute2>=<Value2>...)
        
        Entity Types to Extract:
        1. DRUG - Extract: 通用名, 商品名, 规格, 成分, 禁忌症, 适应症, 不良反应, 半衰期, 代谢途径, 批准文号, 生产企业
        2. DISEASE - Extract: ICD-10, 症状, 并发症, 目标人群, 治疗方案
        3. ENZYME - Extract: 酶名称, 功能, 基因编码
        4. GENE - Extract: 基因名称, 多态性位点, 代谢影响
        5. DOSAGE - Extract: 最大剂量, 服用方式, 联合用药, 特殊人群, 起始剂量
        
        Relationship Format: ("relationship"<|><SOURCE><|><TARGET><|><TYPE><|><STRENGTH>)
        
        Strength Scoring (1-10):
        - 10: Black box warnings, absolute contraindications
        - 8-9: Major interactions, primary metabolic pathways
        - 6-7: Moderate interactions, common adverse reactions
        - 3-5: Minor interactions, theoretical concerns
        - 1-2: Minimal clinical significance
        
        Use {record_delimiter} between each entity/relationship.
        End with <|COMPLETE|>
        
        Drug Manual Text:
        {manualText}
        """;
    
    /**
     * Parse drug manual with structured prompting
     */
    public String parseManualWithStructuredPrompt(String manualText) {
        try {
            log.info("Starting structured drug manual parsing");
            
            // Create prompt with template
            PromptTemplate promptTemplate = new PromptTemplate(EXTRACTION_TEMPLATE);
            String formattedPrompt = promptTemplate.render(
                Map.of("manualText", manualText)
            );
            
            // Call LLM with system and user messages
            String response = chatClient.prompt()
                .system(SYSTEM_PROMPT)
                .user(formattedPrompt)
                .call()
                .content();
            
            log.info("Successfully parsed drug manual with Spring AI");
            return response;
            
        } catch (Exception e) {
            log.error("Error in structured parsing", e);
            throw new RuntimeException("Failed to parse drug manual", e);
        }
    }
    
    /**
     * Parse with custom options
     */
    public String parseWithCustomOptions(String manualText, Double temperature, Integer maxTokens) {
        try {
            log.info("Parsing with custom options - temp: {}, maxTokens: {}", temperature, maxTokens);
            
            var promptBuilder = chatClient.prompt()
                .system(SYSTEM_PROMPT)
                .user(EXTRACTION_TEMPLATE.replace("{manualText}", manualText));
            
            // Note: Spring AI M6 doesn't support dynamic options in this way
            // Options should be configured in application.yml or via ChatModel configuration
            String response = promptBuilder
                .call()
                .content();
            
            return response;
            
        } catch (Exception e) {
            log.error("Error in custom parsing", e);
            throw new RuntimeException("Failed to parse with custom options", e);
        }
    }
    
    /**
     * Stream parsing for large documents
     */
    public void parseManualStreaming(String manualText, StreamCallback callback) {
        try {
            log.info("Starting streaming parse");
            
            chatClient.prompt()
                .system(SYSTEM_PROMPT)
                .user(EXTRACTION_TEMPLATE.replace("{manualText}", manualText))
                .stream()
                .content()
                .subscribe(
                    chunk -> callback.onChunk(chunk),
                    error -> callback.onError(error),
                    () -> callback.onComplete()
                );
                
        } catch (Exception e) {
            log.error("Error in streaming parse", e);
            throw new RuntimeException("Failed to stream parse", e);
        }
    }
    
    /**
     * Callback interface for streaming responses
     */
    public interface StreamCallback {
        void onChunk(String chunk);
        void onError(Throwable error);
        void onComplete();
    }
    
    /**
     * Validate and enhance parsed entities
     */
    public String validateAndEnhanceParsing(String manualText, String initialParsing) {
        String validationPrompt = """
            Review the following entity extraction and:
            1. Ensure all required attributes are present
            2. Validate relationship strengths (1-10)
            3. Add any missing entities or relationships
            4. Correct any format errors
            
            Original Text:
            %s
            
            Initial Parsing:
            %s
            
            Return the corrected and enhanced parsing in the same format.
            """.formatted(manualText, initialParsing);
        
        return chatClient.prompt()
            .system("You are a pharmaceutical data validation expert.")
            .user(validationPrompt)
            .call()
            .content();
    }
}